import torch
import torch.nn as nn

class VGG11(nn.Module):

    def __init__(self, in_channels, out_classes=1000):
        super(VGG11, self).__init__()
        self.in_channels = in_channels
        self.out_classes = out_classes
        # convolutional layers
        self.convLayers = nn.Sequential(
            nn.Conv2d(self.in_channels, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        # fully connected linear layers
        self.linearLayers = nn.Sequential(
            nn.Linear(in_features=7*7*512, out_features=4096),
            nn.ReLU(),
            nn.Dropout2d(0.5),
            nn.Linear(in_features=4096, out_features=4096),
            nn.ReLU(),
            nn.Dropout2d(0.5),
            nn.Linear(in_features=4096, out_features=self.out_classes)
        )

    def forward(self, x):
        # convolution
        x = self.convLayers(x)
        # flatten for fully connected layers
        x = x.reshape(x.shape[0], -1)
        # fully connected
        x = self.linearLayers(x)
        return x

if __name__ == "__main__":
    vgg11 = VGG11(in_channels=3)
    total_params = sum(p.numel() for p in vgg11.parameters())
    print(f"[INFO]: {total_params:,} total parameters.")

    image_tensor = torch.randn(1, 3, 224, 224)
    outputs = vgg11(image_tensor)
    print(outputs.shape)